import { Controller } from '@nestjs/common';
import { Tool } from '@rekog/mcp-nest';
import { z } from 'zod';
import { ToolResponse } from '../type';

@Controller('picture-tool')
export class PictureToolController {
  @Tool({
    name: 'picture-tool',
    description: '画图工具',
    parameters: z
      .object({
        prompt: z
          .string({
            description: '画图的关键词'
          })
          .default('画一只猫')
      })
      .required()
  })
  async exec({
    prompt
  }: {
    // 画图的关键词
    prompt: string;
  }): Promise<ToolResponse> {
    if (!prompt) {
      return {
        isError: true,
        content: [
          {
            type: 'text',
            text: '画图失败'
          }
        ]
      };
    }

    if (!process.env.PICTURE_TOKEN) {
      return {
        isError: true,
        content: [
          {
            type: 'text',
            text: '没有设置PICTURE_TOKEN'
          }
        ]
      };
    }
    if (!process.env.PICTURE_MODEL) {
      return {
        isError: true,
        content: [
          {
            type: 'text',
            text: '没有设置PICTURE_MODEL'
          }
        ]
      };
    }

    let result = await fetch('https://api.siliconflow.cn/v1/images/generations', {
      method: 'POST',
      body: JSON.stringify({
        model: process.env.PICTURE_MODEL,
        prompt: prompt,
        negative_prompt: 'nsfw',
        image_size: '1024x1024',
        batch_size: 1,

        num_inference_steps: 20,
        guidance_scale: 7.5
        // image: 'data:image/webp;base64, XXX',
      }),
      headers: {
        'Content-Type': 'application/json',

        Authorization: 'Bearer ' + process.env.PICTURE_TOKEN
      }
    });

    if (result.status !== 200) {
      return {
        isError: true,
        content: [
          {
            type: 'text',
            text: '画图失败'
          }
        ]
      };
    }

    let images = await result.json();
    if (!images.data) {
      return {
        isError: true,
        content: [
          {
            type: 'text',
            text: '画图失败'
          }
        ]
      };
    }

    return {
      content: [
        {
          type: 'text',
          text: JSON.stringify({
            images: images.data.map((image: any) => image.url)
          })
        }
      ]
    };
  }
}
